/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lvyh.lightframe.tracer.core.common;

import com.lvyh.lightframe.tracer.core.LogData;
import com.lvyh.lightframe.tracer.core.SimpleSpan;
import com.lvyh.lightframe.tracer.core.SimpleSpanBuilder;
import com.lvyh.lightframe.tracer.core.SimpleSpanContext;
import com.lvyh.lightframe.tracer.core.SimpleTracer;
import com.lvyh.lightframe.tracer.core.enums.SamplerType;
import com.lvyh.lightframe.tracer.core.context.SimpleTraceContext;
import com.lvyh.lightframe.tracer.core.context.SimpleTraceContextHolder;
import com.lvyh.lightframe.tracer.core.samplers.PercentageBasedSampler;
import com.lvyh.lightframe.tracer.core.samplers.Sampler;
import com.lvyh.lightframe.tracer.core.samplers.TraceSamplerConfig;
import io.opentracing.tag.Tags;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;


public abstract class AbstractTracer {
    private static final Logger logger = LoggerFactory.getLogger(AbstractTracer.class);

    protected SimpleTracer simpleTracer;

    public AbstractTracer(TraceSamplerConfig traceSamplerConfig) {
        Sampler sampler = null;
        SamplerType samplerType = SamplerType.convert(traceSamplerConfig.getSamplerType());
        if (traceSamplerConfig != null && samplerType == SamplerType.PERCENTAGE_BASED) {
            sampler = new PercentageBasedSampler(traceSamplerConfig.getPercentage());
        } else {
            throw new UnsupportedOperationException("Other sampler type not supported.");
        }
        this.simpleTracer = new SimpleTracer(sampler, null, null);
    }

    public SimpleSpan errorRecover(Map<String, String> baggage) {
        SimpleSpanContext spanContext = SimpleSpanContext.rootStart();
        spanContext.setBaggage(baggage);
        SimpleSpan span = this.genSeverSpanInstance(System.currentTimeMillis(), "", spanContext, null);
        return span;
    }

    protected SimpleSpan genSeverSpanInstance(long startTime, String operationName, SimpleSpanContext simpleSpanContext, Map<String, ?> tags) {
        return new SimpleSpan(this.simpleTracer, startTime, null, operationName, simpleSpanContext, tags);
    }

    /**
     * Stage CS , This stage will produce a new span
     * If there is a span in the current traceContext, it is the parent of the current Span
     *
     * @param operationName as span name
     * @return a new span
     */
    public SimpleSpan clientSend(String operationName) {

        SimpleTraceContext simpleTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();
        SimpleSpan serverSpan = simpleTraceContext.pop();
        SimpleSpan clientSpan = null;
        try {
            // Create a new Span
            SimpleSpanBuilder spanBuilder = (SimpleSpanBuilder) this.simpleTracer.buildSpan(operationName);
            clientSpan = (SimpleSpan) spanBuilder.newChild(serverSpan).start();
            clientSpan.setParentSpan(serverSpan);
            return clientSpan;
        } catch (Throwable throwable) {
            logger.error("client send error and restart by root span", throwable);
            Map<String, String> baggage = null;
            if (serverSpan != null) {
                baggage = serverSpan.getSimpleSpanContext().getBaggage();
            }
            clientSpan = this.errorRecover(baggage);
        } finally {
            if (clientSpan != null) {
                clientSpan.setTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_CLIENT);
                clientSpan.setTag(TraceContants.CURRENT_THREAD_NAME, Thread.currentThread().getName());

                clientSpan.log(LogData.CLIENT_SEND_EVENT_VALUE);

                // Put into the thread context
                clientSpan.setSimpleTracer(simpleTracer);
                simpleTraceContext.add(clientSpan);
            }
        }
        return clientSpan;
    }


    /**
     * Span finished and append tags
     */
    public void clientReceiveTagFinish(SimpleSpan clientSpan, String resultCode) {
        if (clientSpan != null) {
            SimpleTraceContext simpleTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();

            clientSpan.log(LogData.CLIENT_RECV_EVENT_VALUE);
            clientSpan.setTag(TraceContants.RESULT_CODE, resultCode);
            clientSpan.finish();

            if (clientSpan.getParentSpan() != null) {
                simpleTraceContext.add(clientSpan.getParentSpan());
            }
        }
    }

    /**
     * Stage CR, This stage will end a span
     */
    public void clientReceive(String resultCode) {
        SimpleTraceContext simpleTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();

        SimpleSpan clientSpan = simpleTraceContext.pop();
        if (clientSpan == null) {
            return;
        }

        clientSpan.setTag(TraceContants.RESULT_CODE, resultCode);

        // Finish the Span and Report remote server.
        clientSpan.finish();

        // restore parent span
        if (clientSpan.getParentSpan() != null) {
            simpleTraceContext.add(clientSpan.getParentSpan());
        }
    }


    /**
     * server receive request
     */
    public SimpleSpan serverReceive(SimpleSpanContext sofaTracerSpanContext) {
        SimpleSpan newSpan = null;
        SimpleTraceContext sofaTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();
        SimpleSpan serverSpan = sofaTraceContext.pop();
        try {
            if (serverSpan == null) {
                newSpan = (SimpleSpan) this.simpleTracer.buildSpan("")
                        .asChildOf(sofaTracerSpanContext).start();
            } else {
                SimpleSpanBuilder spanBuilder = (SimpleSpanBuilder) this.simpleTracer.buildSpan("");
                newSpan = (SimpleSpan) spanBuilder.newChild(serverSpan).start();
            }
        } catch (Throwable throwable) {
            logger.error("server received and restart root span", throwable);
            Map<String, String> baggage = null;
            if (serverSpan != null) {
                baggage = serverSpan.getSimpleSpanContext().getBaggage();
            }
            newSpan = this.errorRecover(baggage);
        } finally {
            if (newSpan != null) {
                newSpan.log(LogData.SERVER_RECV_EVENT_VALUE);
                newSpan.setTag(Tags.SPAN_KIND.getKey(), Tags.SPAN_KIND_SERVER);
                newSpan.setTag(TraceContants.CURRENT_THREAD_NAME, Thread.currentThread().getName());
                sofaTraceContext.add(newSpan);
            }
        }
        return newSpan;
    }

    /**
     * Stage SS, This stage will end a span
     */
    public void serverSend(String resultCode) {
        try {
            SimpleTraceContext simpleTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();
            SimpleSpan serverSpan = simpleTraceContext.pop();
            if (serverSpan == null) {
                return;
            }
            serverSpan.log(LogData.SERVER_SEND_EVENT_VALUE);
            serverSpan.setTag(TraceContants.RESULT_CODE, resultCode);

            serverSpan.setSimpleTracer(this.simpleTracer);

            Sampler sampler = simpleTracer.getSampler();
            //sampler is support &  current span is root span
            if (sampler != null && serverSpan.getParentSpan() == null) {
                serverSpan.getSimpleSpanContext().setSampled(true);
                serverSpan.getSimpleSpanContext().setSampler(sampler);
            }

            serverSpan.finish();
        } finally {
            this.clearTreadLocalContext();
        }
    }

    /**
     * Clean up all call context information: Note that the server can be cleaned up after receiving it;
     * the client does not have the right time to clean up (can only judge size <= 1)
     */
    private void clearTreadLocalContext() {
        SimpleTraceContext simpleTraceContext = SimpleTraceContextHolder.getSimpleTraceContext();
        simpleTraceContext.remove();
    }

    public SimpleTracer getSimpleTracer() {
        return simpleTracer;
    }

}
